#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <math.h>
#include <string.h>
#include "include/c_api/types_c.h"
#include "include/c_api/context_c.h"
#include "include/c_api/model_c.h"

#define TYPE_NAME(x) _Generic((x), \
    int: "int", \
    float: "float", \
    double: "double", \
    char: "char", \
    int *: "int *", \
    float *: "float *", \
    double *: "double *", \
    char *: "char *", \
    void *: "void *", \
    default: "unknown")

#define PRINT_TYPE(variable) printf("The type of " #variable " is: %s\n", TYPE_NAME(variable))

#define INPUT_SIZE 784 // MNIST 28x28 grayscale image

// 使用已提供的二维数组替代图片读取 label: 7
int16_t mnistData[784] = {  0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0, 255,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0, 255, 127,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0, 255,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0, 255,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0, 255,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0, 255,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0, 255,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0, 255,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0, 255,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0, 255, 127,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0, 255,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0, 255,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0, 255,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0, 172, 255,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0};
// 修改预处理函数，直接使用 mnistData 数组
void Preprocess(float *input_data) {
    for (int i = 0; i < INPUT_SIZE; i++) {
        input_data[i] = mnistData[i] / 255.0f; // 归一化到 [0, 1]
    }
}
void PrintOneTensorShape(MSTensorHandle tensor) {
    // 获取张量的形状数组和维度数量
    size_t shape_size;
    const int64_t *shape = MSTensorGetShape(tensor, &shape_size);
    
    if (shape == NULL) {
        printf("Failed to get tensor shape.\n");
        return;
    }

    // 打印张量的形状信息
    printf("Tensor shape: [");
    for (int i = 0; i < shape_size; ++i) {
        printf("%ld", shape[i]);
        if (i < shape_size - 1) {
            printf(", ");
        }
    }
    printf("]\n");
}
//打印张量shape
int PrintAllTensorShapes(MSTensorHandleArray inputs) {
    for (size_t i = 0; i < inputs.handle_num; ++i) {
        printf("Tensor %zu:\n", i + 1);
        PrintOneTensorShape(inputs.handle_list[i]);
    }
    return 0;
}

void PrintModelDetails(MSModelHandle model) {
    // 打印模型的输入信息
    MSTensorHandleArray inputs = MSModelGetInputs(model);
    printf("----------------------------------------------------------\nModel Inputs (%ld):\n", inputs.handle_num);
    for (size_t i = 0; i < inputs.handle_num; ++i) {
        MSTensorHandle tensor = inputs.handle_list[i];
        printf("  Input %ld:\n", i);
        printf("    Name: %s\n", MSTensorGetName(tensor));
        size_t shape_size;
        const int64_t *shape = MSTensorGetShape(tensor, &shape_size);
        printf("    Shape: [");
        for (int j = 0; j < shape_size; ++j) {
            printf("%ld%s", shape[j], (j < shape_size - 1) ? ", " : "");
        }
        printf("]\n");
        printf("    Data Type: %d\n", MSTensorGetDataType(tensor));
        printf("    Data Size: %ld bytes\n", MSTensorGetDataSize(tensor));
    }

    // 打印模型的输出信息
    MSTensorHandleArray outputs = MSModelGetOutputs(model);
    printf("\nModel Outputs (%ld):\n", outputs.handle_num);
    for (size_t i = 0; i < outputs.handle_num; ++i) {
        MSTensorHandle tensor = outputs.handle_list[i];
        printf("  Output %ld:\n", i);
        printf("    Name: %s\n", MSTensorGetName(tensor));
        size_t shape_size;
        const int64_t *shape = MSTensorGetShape(tensor, &shape_size);
        printf("    Shape: [");
        for (int j = 0; j < shape_size; ++j) {
            printf("%ld%s", shape[j], (j < shape_size - 1) ? ", " : "");
        }
        printf("]\n");
        printf("    Data Type: %d\n", MSTensorGetDataType(tensor));
        printf("    Data Size: %ld bytes\n", MSTensorGetDataSize(tensor));
    }

    // 估算模型大小
    int64_t model_size = 0;
    for (size_t i = 0; i < inputs.handle_num; ++i) {
        model_size += MSTensorGetDataSize(inputs.handle_list[i]);
    }
    for (size_t i = 0; i < outputs.handle_num; ++i) {
        model_size += MSTensorGetDataSize(outputs.handle_list[i]);
    }
    printf("\nEstimated Model Size: %ld bytes\n----------------------------------------------------------\n", model_size);
}

int main(int argc, const char **argv) {
    if (argc < 2) {
        printf("Usage: %s <model.ms>\n", argv[0]);
        return 1;
    }

    const char *model_path = argv[1];

    // 创建并初始化上下文
    MSContextHandle context = MSContextCreate();
    if (context == NULL) {
        printf("MSContextCreate failed.\n");
        return 1;
    }
    MSDeviceInfoHandle cpu_device_info = MSDeviceInfoCreate(kMSDeviceTypeCPU);
    if (cpu_device_info == NULL) {
        printf("MSDeviceInfoCreate failed.\n");
        MSContextDestroy(&context);
        return 1;
    }
    MSDeviceInfoSetEnableFP16(cpu_device_info, false);
    MSContextAddDeviceInfo(context, cpu_device_info);

    // 创建模型
    MSModelHandle model = MSModelCreate();
    if (model == NULL) {
        printf("MSModelCreate failed.\n");
        MSContextDestroy(&context);
        return 1;
    }

    // 从文件构建模型
    int ret = MSModelBuildFromFile(model, model_path, kMSModelTypeMindIR, context);
    if (ret != kMSStatusSuccess) {
        printf("MSModelBuildFromFile failed, ret: %d.\n", ret);
        MSModelDestroy(&model);
        MSContextDestroy(&context);
        return 1;
    }

    //打印模型信息
    PrintModelDetails(model);

    // 获取输入张量
    MSTensorHandleArray inputs = MSModelGetInputs(model);
    if (inputs.handle_list == NULL || inputs.handle_num == 0) {
        printf("MSModelGetInputs failed.\n");
        MSModelDestroy(&model);
        MSContextDestroy(&context);
        return 1;
    }

    //打印张量形状
    printf("inputs: ");
    PrintAllTensorShapes(inputs);

    // 准备输入数据
    float *input_data = (float *)malloc(INPUT_SIZE * sizeof(float));
    if (!input_data) {
        printf("Failed to allocate memory for input data.\n");
        MSModelDestroy(&model);
        MSContextDestroy(&context);
        return 1;
    }
    Preprocess(input_data); // 使用替代的 mnistData 数据

    // 将输入数据拷贝到张量中
    MSTensorHandle input_tensor = inputs.handle_list[0]; //多此一举
    // PRINT_TYPE(inputs.handle_list[0]);
    memcpy(MSTensorGetMutableData(input_tensor), input_data, INPUT_SIZE * sizeof(float));
    free(input_data);

    // 执行推理
    MSTensorHandleArray outputs;
    ret = MSModelPredict(model, inputs, &outputs, NULL, NULL);
    if (ret != kMSStatusSuccess) {
        printf("MSModelPredict failed, ret: %d.\n", ret);
        MSModelDestroy(&model);
        MSContextDestroy(&context);
        return 1;
    }

    //打印张量形状
    printf("outputs: ");
    PrintAllTensorShapes(outputs);

    // 处理输出
    const float *output_data = (const float *)MSTensorGetData(outputs.handle_list[0]);
    int predicted_digit = 0;
    float max_value = output_data[0];
    for (int i = 1; i < 10; i++) {
        if (output_data[i] > max_value) {
            max_value = output_data[i];
            predicted_digit = i;
        }
        printf("outputs[%d]: %.3f\n", i, output_data[i]);
    }

    // 输出结果
    printf("Predicted digit: %d\n", predicted_digit);

    // 清理资源

    if(model != NULL)
    {
        MSModelDestroy(&model);
    }
    return 0;
}
